import asyncio
import json
import os
import random
import re

from typing import List

from agents import (
    Agent,
    ModelSettings,
    RunContextWrapper,
    Runner,
    RunResult,
    function_tool,
    handoff,
    set_tracing_disabled,
)
from agents.models.openai_provider import OpenAIProvider
from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX

from .game_helpers import get_random_card_list, get_cached_expressions

set_tracing_disabled(disabled=True)

old_arrays = []
last_cards_posted = ""

API_KEY = os.getenv("EXAMPLE_API_KEY")
BASE_URL = os.getenv("EXAMPLE_BASE_URL")

provider = OpenAIProvider(
        api_key=API_KEY,
        base_url=BASE_URL,
        use_responses=False
    )
model = provider.get_model("qwen3:8b")
# model = provider.get_model("qwen3-235b-a22b")


instruction_template_dict = {
    "GameDealer": """
        Generate 4 random natural numbers between 1 and 13, include 1 and 13. Just return 4 numbers in an array, don't include other content. The returned array should not be repeated with the following arrays:
        {old_arrays}
        /no_think
        """,
    "MathProdigy": """You are a helpful agent. /no_think""",
    "GameJudger": """You are a helpful agent. /no_think""",
}


user_prompt_template_dict = {
    "GameDealer": """generate an array /no_think""",
    "MathProdigy": """
        What's the 24 points expression of {last_cards_posted} ? 
        If the result is 'expression not found', just return 'expression not found'. 
        If the result is an arithmetic expression, return the expression follow this Python dictionary format:{{"expression": "result"}}, just return this Python dictionary and do not add anything else.
        /no_think
        """,
    "GameJudger": """
        Cards posted is '{last_cards_posted}', what's the check result of {expression} ? Call a function to check whether the expression is correct or wrong, don't resolve it by yourself. Return the check result follow this Python dictionary format:{{"check_result": "result"}}, the result shoule be 'Correct' or 'Wrong', just return this Python dictionary and do not add anything else. 
        /no_think
        """,
}


class CustomContext:
    def __init__(self, context_variables: dict):
        self.context_variables = context_variables


def get_instruction(
    run_context: RunContextWrapper[CustomContext], agent: Agent[CustomContext]
) -> str:
    global instruction_template_dict

    context = run_context.context
    context_variables = context.context_variables

    agent_name = context_variables["agent_name"]
    instruction_template = instruction_template_dict[agent_name]
    instruction = instruction_template

    if agent_name == "GameDealer":
        last_cards_posted = context_variables["old_arrays"]
        instruction = instruction_template.format(old_arrays=last_cards_posted)
   
    return instruction


def get_user_prompt(context_variables):
    global user_prompt_template_dict

    agent_name = context_variables["agent_name"]
    user_prompt_template = user_prompt_template_dict[agent_name]
    user_prompt = user_prompt_template

    if agent_name == "MathProdigy":
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(last_cards_posted=last_cards_posted)

    elif agent_name == "GameJudger":
        expression = context_variables["expression"]
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(expression=expression,last_cards_posted=last_cards_posted)

    elif agent_name == "GamePlayer":
        last_cards_posted = context_variables["last_cards_posted"]
        user_prompt = user_prompt_template.format(last_cards_posted=last_cards_posted)

    return user_prompt


def remove_think_tags(text):
    new_text = re.sub(r"<think>.*?</think>\n?", "", text, flags=re.DOTALL).strip()
    return new_text


@function_tool
def get_24_points_expression_func(last_cards_posted: list) -> str:
    """Resolve the expression of 24 points game, return an arithmetic expression.

    Keyword arguments:
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    point_list = last_cards_posted  
    if len(point_list) == 0:
        return "expression not found"

    expressions = get_cached_expressions(point_list)
    
    result = "expression not found"
    if len(expressions) > 0:
        random_idx = random.randint(0, len(expressions)-1)
        expression = f"'{expressions[random_idx]}'".replace("'", "")
        print(f"The resolved 24 points expression is '{expression}'")
        return expression
    
    return result


@function_tool
def check_24_points_expression_func(expression: str, last_cards_posted: str) -> str:
    """Check if the result of an arithmetic expression is equal 24, return 'Correct' or 'Wrong'.

    Keyword arguments:
      expression: an arithmetic expression
      last_cards_posted: an array of 4 integers between 1 to 13.
    """

    result = eval(expression.replace("'", ""))
    if abs(result - 24) < 0.001:
        return "Correct"

    return "Wrong"


agent_bill = Agent(
    name="GameDealer",
    instructions=get_instruction,
    model=model,
    model_settings=ModelSettings(extra_body={"enable_thinking": False}),
    tools=[],
)

agent_gauss = Agent(
    name="MathProdigy",
    handoff_description="A math prodigy who can help users",    
    instructions=get_instruction,
    model=model,
    model_settings=ModelSettings(extra_body={"enable_thinking": False}),
    tools=[get_24_points_expression_func],
    tool_use_behavior="stop_on_first_tool"
)

agent_peter = Agent(
    name="GameJudger",
    handoff_description="Person responsible for checking if the expression is correct",     
    instructions=get_instruction,
    model=model,
    model_settings=ModelSettings(extra_body={"enable_thinking": False}),
    tools=[check_24_points_expression_func],
    tool_use_behavior="stop_on_first_tool"
)


async def on_handoff_gauss(ctx: RunContextWrapper[None]):
    # print(f"in on_handoff_gauss(), context_variables is {ctx.context.context_variables} ")
    pass

async def on_handoff_peter(ctx: RunContextWrapper[None]):
    # print(f"in on_handoff_peter(), context_variables is {ctx.context.context_variables} ")
    pass

handoff_gauss = handoff(
    agent=agent_gauss,
    on_handoff=on_handoff_gauss,
    tool_name_override="help_agent",
    tool_description_override="an agent who can help user to calculate 24 points expression",
)

handoff_peter = handoff(
    agent=agent_peter,
    on_handoff=on_handoff_peter,
    tool_name_override="judge_agent",
    tool_description_override="an agent who can judge whether the 24 points expression is correct",
)

triage_agent = Agent[CustomContext](
    name="Triage Agent",
    handoff_description="A triage agent that can delegate a customer's request to the appropriate agent.",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX} "
        "You are a helpful triaging agent. You can use your tools to delegate questions to other appropriate agents."
    ),
    model=model,
    handoffs=[
        handoff_gauss,
        handoff_peter,
    ],
    model_settings=ModelSettings(extra_body={"enable_thinking": False})
)


async def deal_cards(old_arrays: List[int]) -> str:
    global agent_bill

    print(f"used old_arrays is :{old_arrays}")
    context_var_dict = {
        "agent_name":"GameDealer", 
        "old_arrays": f"{old_arrays}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_bill,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    cards_posted = remove_think_tags(response.final_output)
    print(cards_posted)
    return cards_posted


async def call_deal():
    global old_arrays, last_cards_posted, agent_bill

    if last_cards_posted != "":
        old_arrays.append(json.loads(last_cards_posted))

    print(f"used old_arrays is :{old_arrays}")
    context_var_dict = {
        "agent_name":"GameDealer", 
        "old_arrays": f"{old_arrays}"
    }
    context = CustomContext(context_variables=context_var_dict)
    response = await Runner.run(
        agent_bill,
        [{"role": "user", "content": get_user_prompt(context_var_dict)}],
        context=context
    )

    last_cards_posted = remove_think_tags(response.final_output)
    print(last_cards_posted)
    cards_posted = str(last_cards_posted)[1:-1]
    return cards_posted


async def call_help():
    global last_cards_posted
    
    context_var_dict = {
        "agent_name":"MathProdigy", 
        "last_cards_posted": json.loads(last_cards_posted)
    }
    context = CustomContext(context_variables=context_var_dict)
    
    user_input = "I want to talk to agent MathProdigy /no_think"
    result: RunResult = await Runner.run(triage_agent, user_input, context=context)
    if result and result.last_agent in [agent_gauss]:
        print(f'Handoff to {result.last_agent.name}')
        sub_result: RunResult = await Runner.run(
            result.last_agent,
            [{"role": "user", "content": get_user_prompt(context_var_dict)}],
            context=context
        )
        print(f"Assistant [{result.last_agent.name}] => {sub_result.final_output}")
        llm_resp = sub_result.final_output
        return llm_resp 
    
    return result.final_output


async def call_judge(expression: str):
    global last_cards_posted

    context_var_dict = {
        "agent_name":"GameJudger", 
        "expression": f"{expression}",
        "last_cards_posted": last_cards_posted
    }

    context = CustomContext(context_variables=context_var_dict)

    user_input = "I want to talk to agent GameJudger /no_think"
    result: RunResult = await Runner.run(triage_agent, user_input, context=context)
    if result and result.last_agent in [agent_peter]:
        print(f'Handoff to {result.last_agent.name}')
        sub_result: RunResult = await Runner.run(
            result.last_agent,
            [{"role": "user", "content": get_user_prompt(context_var_dict)}],
            context=context
        )
        print(f"Assistant [{result.last_agent.name}] => {sub_result.final_output}")
        llm_resp = remove_think_tags(sub_result.final_output)
        return llm_resp 
    
    return result.final_output
